import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import os

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)


class ConvCell(keras.Model):

    def __init__(self, filters, ksize=(3, 3), strides=(1, 1), padding='same', **kwargs):
        super().__init__(**kwargs)
        self.conv = layers.Conv2D(filters, ksize, strides, padding, use_bias=False)
        self.bn = layers.BatchNormalization()
        self.relu = layers.LeakyReLU()

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.bn(x, training=training)
        x = self.relu(x, training=training)
        return x


class Residual(keras.Model):

    def __init__(self, filters, n_residual, **kwargs):
        super().__init__(**kwargs)
        self.residuals = []
        for i in range(n_residual):
            self.residuals.append(keras.Sequential([
                ConvCell(filters // 2, (1, 1)),
                ConvCell(filters, (3, 3))
            ]))

    def call(self, inputs, training=None, mask=None):
        x = inputs
        for i, residual in enumerate(self.residuals):
            x = residual(x) + x
        return x


class Darknet53(keras.Model):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = ConvCell(32)
        self.cfg = [
            [64, 1],
            [128, 2],
            [256, 8],
            [512, 8],
            [1024, 4]
        ]
        self.convs = []
        for i, item in enumerate(self.cfg):
            filters = item[0]
            n_residual = item[1]
            self.convs.append(ConvCell(filters, (3, 3), (2, 2)))
            self.convs.append(Residual(filters, n_residual))
        self.avg = layers.GlobalAvgPool2D()
        self.fc = layers.Dense(n_cls, activation=activations.softmax)

    def call(self, inputs, training=None, mask=None):
        print('inputs', inputs.shape)
        x = self.conv1(inputs)
        print('after conv', x.shape)
        for i, layer in enumerate(self.convs):
            x = layer(x)
            print(f'#{i}: {x.shape}')
        x = self.avg(x)
        print('after g avg', x.shape)
        x = self.fc(x)
        print('after fc', x.shape)
        return x


if '__main__' == __name__:
    model = Darknet53(10)
    model.build(input_shape=(None, 256, 256, 3))
    model.summary()

    x = tf.zeros((4, 256, 256, 3), dtype=tf.float32)
    pred = model(x)
    print(pred.shape)
